"""
Created on 01-Mar-2010

@author: dexen
"""


import time

import rpyc


from dexen.system.common.base import Action
from dexen.system.common.service import DexenAddr
from dexen.system.common.states import SlaveState, M_CRASHED, M_RUNNING, M_FINISHED, M_STOPPED
from dexen.system.common.transact import SlaveTransact, SlaveRecvTransact
from dexen.system.common.meta import JobMeta
from dexen.system.common.excepts import JobAlreadyStoppedException,\
    JobIsRunningException

#===============================================================================
# Constants
#===============================================================================
N_RUNNING = "NODE RUNNING"
N_CRASHED = "NODE CRASHED"


INF = 100000
MAX_SLAVE_SLOT = 2

#===============================================================================
# Structural Classes
#===============================================================================
class NodeProxy(object):
    def __init__(self, node_name, node_addr, reg_addr):
        assert isinstance(node_addr, DexenAddr)
        assert isinstance(reg_addr, DexenAddr)
        self.name = node_name
        self.masters = {} # job_id:MasterProxy
        self.slaves = {} # slave_name:SlaveState 
        self._state = N_RUNNING
        self.reg_addr = reg_addr
        try:
            self.conn = node_addr.get_conn()
            self.num_master_slots = self.conn.root.n_master_slots()
            self.num_master_slots_taken = 0     
        except EOFError:
            self._state = N_CRASHED
    
    def get_master(self, job_id):
        return self.masters[job_id]
    
    def create_master(self, master_name, job_meta):
        try:   
            self.conn.root.create_master(master_name) # asking the node to create the master
        except EOFError:
            self._state = N_CRASHED
            return None
        
        conn = self.reg_addr.get_conn()
        for _ in range(10):
            time.sleep(0.5)        
            master_addr = conn.root.get_master_addr(master_name)
            master_addr = rpyc.classic.obtain(master_addr)
            if master_addr != None: break
        conn.close()
        if not master_addr: pass #TODO:
        assert isinstance(job_meta, JobMeta)
        self.masters[job_meta.get_id()] = MasterProxy(master_name, master_addr, job_meta)
    
    def start_master(self, job_id):
        master = self.masters[job_id]
        assert isinstance(master, MasterProxy)
        master.run_job()
        self.num_master_slots_taken += 1
    
    def stop_master(self, job_id, state):
        master = self.masters[job_id]
        assert isinstance(master, MasterProxy)
        master.stop_job(state)
        self.num_master_slots_taken -= 1
            
    def del_master(self, job_id):
        master = self.masters[job_id]
        assert isinstance(master, MasterProxy)
        if master.is_running():
            raise JobIsRunningException("Cannot delete a running job.")
        del self.masters[job_id]
    
    def n_slave_slots(self):
        try:
            return self.conn.root.n_slave_slots()
        except:
            self._state = N_CRASHED
            return 0
        
    def create_slave(self, name):
        try:
            self.conn.root.create_slave(name) # asking the node to create the slave
        except:
            self._state = N_CRASHED
            return None
        slave_state = SlaveState(name)
        self.slaves[name] = slave_state
    
    def get_slave(self, name):
        return self.slaves[name]
    
    def get_free_slave(self):
        for slave in self.slaves.values():
            assert isinstance(slave, SlaveState)
            if slave.is_idle():
                return slave
        return None
    
    def has_master_slot(self):
        return self.num_master_slots_taken < self.num_master_slots

    def ping(self):
        try:
            self.conn.root.is_alive()
        except EOFError:
            self._state = N_CRASHED
    
    def is_running(self):
        return self._state == N_RUNNING
    
    def is_crashed(self):
        return self._state == N_CRASHED
    
    def __repr__(self):
        ret = ""
        for name in self.slaves:
            ret += "    %s\n" %(name)
        for name in self.masters:
            ret += "    %s\n" % (name)
        return ret
        
class MasterProxy(object):
    def __init__(self, master_name, master_addr, job_meta):
        self.name = master_name
        assert isinstance(master_addr, DexenAddr)
        self.conn = master_addr.get_conn()
        self.state = None
        self.scheduled_slaves = set()
        self.job_meta = job_meta
        self.new_actions = []
        
    def run_job(self):
        try:
            self.conn.root.run_job(self.job_meta)
            self.state = M_RUNNING
        except EOFError:
            self.state = M_CRASHED

    def stop_job(self, state):
        try:
            if self.state != M_RUNNING:
                raise JobAlreadyStoppedException("Job is already stopped.")
            self.conn.root.stop_job()
        except EOFError:
            pass
        self.state = state

    def del_job(self):
        pass

    def is_crashed(self):
        return self.state == M_CRASHED
    
    def on_crash(self):
        self.scheduled_slaves = set()
        
    def is_running(self):
        return self.state == M_RUNNING
    
    def get_state(self):
        return self.state

    def get_actions(self):
        actions = self.new_actions
        self.new_actions = []
        return actions
    
    def update(self):
        """
        During getting the updated state from master, if the connection is broken
        Notice that the self.state contains the latest state before it crashes
        Note that this proxy object has to turn on the flag crashed in the state
        object
        """
        try:
            actions = self.conn.root.get_updates()
            actions = rpyc.classic.obtain(actions)
            for action in actions:
                assert isinstance(action, Action)
                if action.is_slave():
                    transact = action.get_transact()
                    assert isinstance(transact, SlaveTransact)
                    if isinstance(transact, SlaveRecvTransact):
                        self.scheduled_slaves.remove(transact.slave_name)
            self.new_actions.extend(actions)
        except EOFError:
            self.state = M_CRASHED
    
    def can_schedule(self):
        try:
            return self.conn.root.can_schedule()
        except EOFError:
            self.state = M_CRASHED
            return False
        
    def schedule(self, slave_name):
        try:
            self.scheduled_slaves.add(slave_name)
            self.conn.root.schedule(slave_name)
        except EOFError:
            self.state = M_CRASHED

    def get_scheduled_slaves(self):
        return self.scheduled_slaves
    
    def reset_scheduled_slaves(self):
        self.scheduled_slaves.clear()  
    